from __future__ import print_function
import sys

# This is not required if you've installed pycparser into
# your site-packages/ with setup.py
#
sys.path.insert(0, '..')

from pycparser import c_parser, c_ast, parse_file
from copy import *


def is_in(item,collection):
	typ=type(item)

	for i in collection:
		if(item.equals(i[1])):
			return True
	return False

def GetFieldFromStructDef(field_name,struct_def):
	result = None
	for f in struct_def:
		if field_name == f.name:
			result = f
			break
	if result == None:
		print('Unable to find field '+field_name+' in struct')
		exit()
	return result
		
class Function(object):
	def __init__(self,set_name):
		self.writes=[]
		self.reads=[]
		self.possible_reads=[]
		self.possible_writes=[]
		self.redirections=[]
		self.name=set_name
		self.locals={}
		self.params={}
		self.params_list=[]
		self.calls=[]
		self.called_by=[]
		self.predecessors=[]
		self.last_inserted=None
	def __str__(self):
		return '{ '+str(self.reads)+' -> (('+self.name+')) -> '+str(self.writes)+' }'
	def isempty(self):
		return len(self.writes) == 0 and len(self.reads) == 0 and len(self.possible_reads) == 0 and len(self.possible_writes) == 0

class Statement(object):
	def __init__(self,enclosing_function):
		self.prev = enclosing_function.last_inserted
		enclosing_function.last_inserted = self
	def equals(self,other):
		return self.var.equals(other.var)

class Call(Statement):
	def __init__(self,func, enclosing_function):
		super(Call,self).__init__(enclosing_function)
		self.func = func

class Write(Statement):
	def __init__(self,var,enclosing_function):
		super(Write,self).__init__(enclosing_function)
		self.var = var

class Possible(Statement):
	def __init__(self,var,enclosing_function):
		super(Possible,self).__init__(enclosing_function)
		self.var = var

class Read(Statement):
	def __init__(self,var,enclosing_function):
		super(Read,self).__init__(enclosing_function)
		self.var = var

class FunctionAnalyzer(object):
	""" Uses the same visitor pattern as c_ast.NodeVisitor, but modified to
	return a value from each visit method, using string accumulation in 
	generic_visit.
	"""
	def __init__(self,set_functions,set_globals,set_struct_defs):
		self.globals=set_globals
		self.functions=set_functions
		if(not ('___global___' in self.functions)):
			self.functions['___global___']=Function('___global___')
		self.struct_defs = set_struct_defs
		self.current_function=None

    
	def visit(self, node):
		method = 'visit_' + node.__class__.__name__
		return getattr(self, method, self.generic_visit)(node)
    
	def generic_visit(self, node):
		if (not (node is None)):
			for c in node.children():
				self.visit(c)

	def AddFuncToCalls(self,foo):
		func = self.current_function
		if(foo in self.functions):
			self.functions[foo].called_by.append(Call(self.functions[func],self.functions[func]))
			self.functions[func].calls.append(self.functions[foo])

	def AddVarToReads(self,var):
		func = self.current_function
		if(func == None):
			func = '___global___'
		self.AddToList(Read(var,self.functions[func]), self.functions[func].reads)

	def AddVarToWrites(self,var):
		func = self.current_function
		if(func == None):
			func = '___global___'
		self.AddToList(Write(var,self.functions[func]), self.functions[func].writes)

	def AddVarToPossibleWrites(self,var):
		func = self.current_function
		if(func == None):
			func = '___global___'
		self.AddToList(Possible(var,self.functions[func]), self.functions[func].possible_writes)

	def AddToList(self,var,l):
		if(var.var != None and var.var.parent != None):
			scope = self.resolveScope(var.var.var.name)
			if(not is_in(var,l)):
				l.append([ scope, var ])

	def resolveScope(self,var_name):
		scope=None
		if(var_name in self.globals):
			scope='global'
		if(var_name in self.functions[self.current_function].params):
			scope='parameter'
		if(var_name in self.functions[self.current_function].locals):
			scope='local'
		if(scope==None):
			print('Error, unable to find variable '+str(var_name)+' in scope of '+str(self.current_function))
			exit()
		return scope

	def findDef(self, var_name):
		definition = None
		if(var_name in self.globals):
			definition = self.globals[var_name]
		if(var_name in self.functions[self.current_function].params):
			definition = self.functions[self.current_function].params[var_name]
		if(var_name in self.functions[self.current_function].locals):
			definition = self.functions[self.current_function].locals[var_name]
		if(definition==None):
			print('Error, unable to find definition for '+type(var_name)+' in scope of '+str(self.current_function))
			exit()
		return definition

	def visit_Constant(self, n):
		return None
        
	def visit_ID(self, n):
		func=self.functions[self.current_function]
		var = None

		if(n.name in func.locals):
			var = func.locals[n.name]
		elif(n.name in self.globals):
			var = self.globals[n.name]
		elif( n.name in func.params):
			var = func.params[n.name]
		if(var == None):
			return None
		return var

	def visit_ArrayRef(self, n):
		#TODO say the tmp_name is ptr+val. need to know which of those to append the PtrRef(indirection) to

		sub = self.visit(n.subscript)
		if(sub != None):
			self.AddVarToReads(sub)

		tmp_name=n.name
#		indirection=0
#		while(type(tmp_name)==c_ast.ArrayRef):
#			tmp_name=tmp_name.name
#			indirection = indirection + 1

		var=self.visit(tmp_name)
		assert(var != None)
		assert(len(var.children) == 1)
		var = var.children[0]
		return var


	def visit_StructRef(self, n):
		#TODO unnecessary?
		self.visit(n.field)

		ret = self.visit(n.name)
		assert(ret != None)
		assert(ret.type.base_type in self.struct_defs)
		field_name = n.field.name
		for child in ret.children:
#			assert(type(child) == depends.StructDef)
			if(child.field_name == field_name):
				return child

		print('Error finding field '+field_name+' in struct '+ret[0].type.base_type)
		exit()

	
	#TODO function pointers? probably just assume we have no idea where they are going
	def visit_FuncCall(self, n):
		if(n.args != None):
			for arg in n.args.exprs:
				reference=self.visit(arg)
				if(reference == None): 
					continue
				if(not (n.name.name in self.functions)):
					if(reference.is_ref()):
						assert(len(reference.children) == 1)
						self.AddVarToPossibleWrites(reference.children[0])

				self.AddVarToReads(reference)
		# Could be a function pointer or something? In which case it would reference something and should be included in reads
		self.AddFuncToCalls(n.name.name)
		self.visit(n.name)

	#TODO there are some special cases for unary ops. For instance, & is a binary op and doesn't read anything but ++ is a binary op that both reads and writes
	# The above TODO may apply to Binary Ops as well?
	def visit_UnaryOp(self, n):

		var = self.visit(n.expr)
		# (ptr+3)++ and ++(ptr+3) is considered invalid lvalue for gcc
		if(n.op == '++' or n.op == '--' or n.op == 'p++' or n.op == 'p--'):
			self.AddVarToWrites(var)
			self.AddVarToReads(var)
		elif(n.op == '*'):
			assert(len(var.children) == 1)
			self.AddVarToReads(var)
			var = var.children[0]
		elif(n.op == '&'):
			assert(var.parent != None)
			var = var.parent
		elif(n.op=='+' or n.op == '-' or n.op == '~' or n.op == '!'):
			pass
		elif(n.op=='sizeof'):
			pass
		else:
			print('unhandled unary operator '+n.op)
			exit()
		return var

	def visit_BinaryOp(self, n):
		lvar = self.visit(n.left)
		rvar = self.visit(n.right)
		self.AddVarToReads(lvar)
		self.AddVarToReads(rvar)
		
		assert(not((rvar != None and rvar.is_ref()) and (lvar != None and lvar.is_ref())))
		if(rvar != None and rvar.is_ref()):
			return rvar
		if(lvar != None and lvar.is_ref()):
			return lvar
		return None
    
	def visit_Assignment(self, n):
		rvar = self.visit(n.rvalue)
		lvar = self.visit(n.lvalue)

		self.AddVarToReads(rvar)
		self.AddVarToWrites(lvar)
		if(n.op != '='):
			# for when the operator if +=, -=, *=, etc
			self.AddVarToReads(lvar)

    
	def visit_IdentifierType(self, n):
		return

	def recur_on_expr_list(self, definition):
		if(len(definition.children) == 0):
			return [ definition ]
		else:
			result = [ ]
			for d in definition.children:
				result = result + self.recur_on_expr_list(d)
			return result
    
	def visit_Decl(self, n, no_type=False):
		typ = type(n.type)

		#TODO variable name obfuscation
		if(not no_type):
			var = self._generate_type(n)
			
			if(n.init):
#				print('WARNING: The variable '+n.name+' is being initialized in its declaration, but this isn\'t fully supported yet so the output may not be entirely correct')
				definition = self.findDef(n.name)
				if(type(n.init) == c_ast.ExprList):
		
					written = self.recur_on_expr_list(definition)
					for w in written:
						self.AddVarToWrites(w)
				else:
					self.AddVarToWrites(definition)
		if n.bitsize: self.visit(n.bitsize)
		if n.init:
			var = self.visit(n.init)
			self.AddVarToReads(var)
    
	def visit_DeclList(self, n):
		self.visit(n.decls[0])
		if len(n.decls) > 1:
			for decl in n.decls[1:]:
				self.visit_Decl(decl, no_type=True) 
    
	def visit_Typedef(self, n):
		self._generate_type(n.type)
    
	def visit_Cast(self, n):
		self._generate_type(n.to_type)
		self.visit(n.expr)
   
	# is this only called when initializing an array? 
	def visit_ExprList(self, n):
		for expr in n.exprs:
			var = self.visit(expr)
			self.AddVarToReads(var)
    
	def visit_Enum(self, n):
		if n.values:
			for i, enumerator in enumerate(n.values.enumerators):
				if enumerator.value: 
					self.visit(enumerator.value)
    
	def visit_Struct(self, n):
		if n.decls:
			for decl in n.decls:
				self.visit(decl)
        
	def visit_FuncDef(self, n):
	
		self.current_function=n.decl.name
		self.visit(n.decl)
		# The body is a Compound node
		self.visit(n.body)
		self.current_function=None

	def visit_FileAST(self, n):
		for ext in n.ext:
			self.visit(ext)

	def visit_Compound(self, n):
		if(n.block_items != None):
			for stmt in n.block_items:
				self.visit(stmt)
    
	def visit_ParamList(self, n):
		for param in n.params:
			self.visit(param)

	def visit_Return(self, n):
		var = None
		if n.expr: var = self.visit(n.expr)
		self.AddVarToReads(var)

	def visit_Break(self, n):
		return None
        
	def visit_Continue(self, n):
		return None
    
	def visit_TernaryOp(self, n):
		var = self.visit(n.cond)
		self.AddVarToReads(var)
		var = self.visit(n.iftrue)
		self.AddVarToReads(var)
		var = self.visit(n.iffalse)
		self.AddVarToReads(var)
    
	def visit_If(self, n):
		if n.cond: 
			var = self.visit(n.cond)
			self.AddVarToReads(var)
			
		self.visit(n.iftrue)
		if n.iffalse: 
			self.visit(n.iffalse)
    
	def visit_For(self, n):
		if n.init: self.visit(n.init)
		if n.cond: 
			var = self.visit(n.cond)
			self.AddVarToReads(var)
		if n.next: self.visit(n.next)
		self.visit(n.stmt)

	def visit_While(self, n):
		if n.cond: 
			var = self.visit(n.cond)
			self.AddVarToReads(var)
		self.visit(n.stmt)

	def visit_DoWhile(self, n):
		self.visit(n.stmt)
		if n.cond: 
			var = self.visit(n.cond)
			self.AddVarToReads(var)

	def visit_Switch(self, n):
		var = self.visit(n.cond)
		self.AddVarToReads(var)
		self.visit(n.stmt)
    
	def visit_Case(self, n):
		var = self.visit(n.expr)
		# case expressions are constant but it can't really hurt to try to add them. Maybe a constant declared variable would be useful to know about
		# probably not but it doesn't really hurt to do this anyway. We can set a flag later to tell it not to
		self.AddVarToReads(var)
		self.visit(n.stmt)
    
	def visit_Default(self, n):
		self.visit(n.stmt)

	def visit_Label(self, n):
		self.visit(n.stmt)

	def visit_Goto(self, n):
		return

	def visit_EllipsisParam(self, n):
		return

	def visit_Struct(self, n):
		self._generate_struct_union(n, 'struct')

	def visit_Typename(self, n):
		self._generate_type(n.type)
        
	def visit_Union(self, n):
		self._generate_struct_union(n, 'union')

	def visit_NamedInitializer(self, n):
		self.visit(n.expr)

	def _generate_struct_union(self, n, name):
		if n.decls:
			for decl in n.decls:
				self.visit(decl)

	def _generate_type(self, n, modifiers=[]):
		""" Recursive generation from a type node. n is the type node. 
		modifiers collects the PtrDecl, ArrayDecl and FuncDecl modifiers 
		encountered on the way down to a TypeDecl, to allow proper
		generation from it.
		"""
		typ = type(n)
		
        #~ print(n, modifiers)
        
		if typ == c_ast.TypeDecl:
			self.visit(n.type)
            
            # Resolve modifiers.
            # Wrap in parens to distinguish pointer to array and pointer to
            # function syntax.
            #
			for i, modifier in enumerate(modifiers):
				if isinstance(modifier, c_ast.ArrayDecl):
					self.visit(modifier.dim)
				elif isinstance(modifier, c_ast.FuncDecl):
					self.visit(modifier.args)
			return
		elif typ == c_ast.Decl:
			self._generate_type(n.type)
			return
		elif typ == c_ast.Typename:
			self._generate_type(n.type)
			return
		elif typ in (c_ast.ArrayDecl, c_ast.PtrDecl, c_ast.FuncDecl):
			self._generate_type(n.type, modifiers + [n])
			return
		else:
			self.visit(n)
			return
